import jieba
import jieba.posseg as pseg
import pandas as pd
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
from py2neo import *


jieba.load_userdict("./data/userdict.txt")

import sys, io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")

class QuestionPrediction():
    """
    问题预测: 采用NB进行多分类
    数据集：data/template_train.csv
    """
    def __init__(self):
        # 连接neo4j
        self.graph = Graph("http://localhost:7474",auth=("neo4j","neo4j"))

        # 编号和类别
        self.all_class = {
            0:'属于',
            1:'描述',
            2:'声明',
            3:'返回',
            4:'例子',
            5:'函数'
        }


        # 训练模板数据
        self.train_file = "./data/template_train.csv"


        # 词语重要度构建器
        self.tv = TfidfVectorizer()


        # 读取训练数据
        self.train_x, self.train_y = self.read_train_data(self.train_file)

        
        # 要训练的训练模型
        self.model = self.train_model_NB(self.train_x, self.train_y)


    # 获取训练数据
    def read_train_data(self, template_train_file):
        train_x = []
        train_y = []
        train_data = pd.read_csv(template_train_file)
        train_x = train_data["text"].apply(lambda x: " ".join(list(jieba.cut(str(x))))).tolist()
        train_y = train_data["label"].tolist()
        return train_x,train_y


    # 采用NB训练模板数据
    def train_model_NB(self, X_train, y_train):
        train_data = self.tv.fit_transform(X_train).toarray()
        clf = MultinomialNB(alpha=0.01)
        clf.fit(train_data, y_train)
        return clf

    # 问题预测
    def predict(self,question):
        res = pseg.cut(question)
        target = None
        for word, flag in res:
            print(word+' '+flag)
            if(flag == 'nm'):
                target = word
                break

        if not target:
            return 'no'
        # print("关键词是:",target)
        question = [" ".join(list(jieba.cut(question)))] #sk库设计的就是能处理好多条输入,但这里我们只需要一个
        test_data = self.tv.transform(question).toarray() #把这句话变成向量

        
        # 得到它是第几类问题
        class_num = self.model.predict(test_data)[0]
        class_name = self.all_class[class_num]
        
        print('问题的类别名字是:',class_name)

        # 这里解决的是  函数 -> xxx  的问句
        cql = "match (a:function_name)-[r:{}]->(ans) where a.name ={} return ans".format(class_name,'~\"'+target+'.*\"')
        
        print('the cql is:',cql)
        res = self.graph.run(cql)

        return res.data()[0]['ans']['name']



if __name__ == '__main__':

    question_model = QuestionPrediction()

    res = question_model.predict("栈返回什么")
    print(res)

    
            


